import asyncio
import concurrent.futures
import datetime
import logging
import sys
import subprocess
import time
import typing
from collections import defaultdict
from dataclasses import dataclass

from tornado.ioloop import IOLoop, PeriodicCallback

from coin.base.timestamp import get_timestamp
from coin.util.queue.tools.kafka_topic import parse_kafka_topic


def _get_strats_from_topic(topics: typing.List[str]) -> typing.List[str]:
  return list(set([parse_kafka_topic(topic)[1].strategy_name for topic in topics]))


def _to_td_str(ts):
  td = datetime.datetime.utcfromtimestamp(ts / 10**9)
  td_str = td.date().strftime('%Y%m%d')
  return td_str


@dataclass
class MonitorProcessInfo:
  strats: typing.List[str]
  process: subprocess.Popen


class StrategyStreamProcessorManager(object):
  def __init__(
      self,
      *,
      ioloop: IOLoop,
      exit_after_min: typing.Optional[int],
      topic_groups: typing.List[typing.List[str]],):
    self._start_ts = get_timestamp()
    self._td_str = _to_td_str(self._start_ts)
    self._ioloop = ioloop
    self._exit_after_min = exit_after_min
    all_topics = []
    for topics in topic_groups:
      all_topics.extend(topics)
    self._all_topics = all_topics
    self._strat_groups = [_get_strats_from_topic(topics) for topics in topic_groups]
    tracked_strategies = []
    for strats in self._strat_groups:
      tracked_strategies.extend(strats)
    self._tracked_strategies = list(set(tracked_strategies))
    self._monitor_info: typing.Dict[int, MonitorProcessInfo] = {}
    self._strat_monitor_id_dict: typing.Dict[str, int] = {}
    self._new_monitor_id = None
    self._exit = False
    self._start()

  def gen_sub_monitor_cmd(self, strats: typing.List[str]) -> str:
    raise NotImplementedError()

  def handle_topics(self, all_topics: typing.List[str]):
    logging.info('Check new topics')
    new_topics = set(all_topics).difference(self._all_topics)
    new_strat_topics = set([
        topic for topic in new_topics if parse_kafka_topic(topic)[1].strategy_name
        not in self._strat_monitor_id_dict])
    old_strat_topics = new_topics.difference(new_strat_topics)
    self._handle_old_strat_topic(list(old_strat_topics))
    self._handle_new_strat_topic(list(new_strat_topics))
    self._all_topics = all_topics
    self._tracked_strategies = list(set(_get_strats_from_topic(self._all_topics)))

  def signal_handler(self, sig, frame):
    logging.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self._on_exit)

  def _update_monitor_info(self,
                           monitor_id: int,
                           strats: typing.List[str],
                           process: subprocess.Popen):
    self._monitor_info[monitor_id] = MonitorProcessInfo(
        strats=strats, process=process)

  def _update_strat_monitor_id_dict(self,
                                    strats: typing.List[str],
                                    monitor_id: int):
    for strat in strats:
      self._strat_monitor_id_dict[strat] = monitor_id

  def _kill_monitor(self, monitor_id: int):
    self._monitor_info[monitor_id].process.terminate()

  def _start(self):
    PeriodicCallback(self._check_running_status, 60*1000).start()
    PeriodicCallback(self._check_exit, 10*1000).start()
    for monitor_id, strats in enumerate(self._strat_groups):
      assert monitor_id not in self._monitor_info, monitor_id
      for strat in strats:
        assert strat not in self._strat_monitor_id_dict, strat
      self._launch_sub_monitor(monitor_id, strats)
      time.sleep(0.1)

  def _launch_sub_monitor(self, monitor_id: int, strats: typing.List[str]):
    assert len(strats) > 0
    if monitor_id in self._monitor_info:
      self._kill_monitor(monitor_id)
    cmd = self.gen_sub_monitor_cmd(strats)
    monitor_process = subprocess.Popen(cmd, shell=True)
    self._update_monitor_info(monitor_id, strats, monitor_process)
    self._update_strat_monitor_id_dict(strats, monitor_id)

  def _handle_new_strat_topic(self, topics: typing.List[str]):
    if len(topics) == 0:
      return
    new_strats = _get_strats_from_topic(topics)
    new_worker_needed = True
    if self._new_monitor_id is not None:
      strat_num = len(self._monitor_info[self._new_monitor_id].strats)
      avg_strats_per_worker = len(self._tracked_strategies) / len(self._monitor_info)
      if strat_num < avg_strats_per_worker:
        new_worker_needed = False
    strats = new_strats
    if new_worker_needed:
      monitor_id = max(self._monitor_info.keys()) + 1
      logging.info('Launch new monitor process for new strat: %s' % new_strats)
    else:
      logging.info('Reuse existing monitor process for new strats: %s' % new_strats)
      assert self._new_monitor_id is not None
      monitor_id = self._new_monitor_id
      strats += self._monitor_info[monitor_id].strats
    self._new_monitor_id = monitor_id
    self._launch_sub_monitor(monitor_id, strats)

  def _handle_old_strat_topic(self, topics: typing.List[str]):
    if len(topics) == 0:
      return
    old_strats = _get_strats_from_topic(topics)
    monitor_id_map = defaultdict(list)
    for strat in old_strats:
      monitor_id = self._strat_monitor_id_dict[strat]
      monitor_id_map[monitor_id].append(strat)
    for monitor_id, strats in monitor_id_map.items():
      logging.info('Reuse existing monitor process for old strats: %s' % strats)
      self._launch_sub_monitor(monitor_id, self._monitor_info[monitor_id].strats)

  def _check_running_status(self):
    if self._exit:
      return
    logging.info('Check running status')
    exit_monitor_ids = [
        monitor_id for monitor_id, info in self._monitor_info.items()
        if info.process.poll() is not None]
    for monitor_id in exit_monitor_ids:
      strats = self._monitor_info[monitor_id].strats
      logging.info('Relaunch existing monitor process due to exit: %s' % strats)
      self._launch_sub_monitor(monitor_id, strats)

  def _check_exit(self):
    logging.info('Check exit')
    subprocess_exit = all([info.process.poll() is not None for info in
                           self._monitor_info.values()])
    if subprocess_exit:
      self._exit = True
      logging.info('Exit monitor runner due to all subprocess exit.')
    if (self._exit_after_min is not None and
        get_timestamp() - self._start_ts > self._exit_after_min * 60 * 10**9):
      self._exit = True
      logging.info(
          'Exit monitor runner due to exit after %s min.' % self._exit_after_min)
    td_str = _to_td_str(get_timestamp())
    if td_str != self._td_str:
      self._exit = True
      logging.info('Exit monitor runner due to trading date change.')
    if self._exit:
      self._on_exit()

  def _on_exit(self):
    self._clean_up()
    sys.exit(0)

  def _clean_up(self):
    for monitor_id, info in self._monitor_info.items():
      logging.info(f'terminate monitor_id: {monitor_id}, process_id: {info.process.pid}')
      info.process.terminate()
